import os
import shutil
import tempfile
import uuid

import pandas as pd
import pymysql
from sqlalchemy import create_engine

from configs.db_config import MysqlConfig, DbConfig


class DbConfigUtils:

    @classmethod
    def get_connect(cls, _database_name, _config):
        # 连接数据库
        if isinstance(_config, MysqlConfig):
            return cls.get_mysql_connect(_database_name, _config)
        else:
            raise Exception('不支持的数据库连接类型')

    @classmethod
    def get_mysql_connect(cls, _database_name, _config):
        # 连接数据库
        cnx = pymysql.connect(host=_config.URL,
                              port=_config.PORT,
                              user=_config.USERNAME,
                              password=_config.PASSWORD,
                              database=_database_name)
        return cnx

    @classmethod
    def execute_no_result_sql(cls, _database_name, _sql, _config):
        with DbConfigUtils.get_connect(_database_name, _config) as cnx:
            with cnx.cursor() as cursor:
                for _part_sql in _sql.split(';'):
                    _part_sql = _part_sql.strip()
                    if len(_part_sql) <= 0:
                        continue
                    cursor.execute(_part_sql)
                cnx.commit()

    @classmethod
    def get_table_data(cls, _database_name, _sql, _config, **kwargs):
        with DbConfigUtils.get_connect(_database_name, _config) as cnx:
            df = pd.read_sql(_sql, cnx, **kwargs)
            return df

    @classmethod
    def insert_table_data(cls, _database_name, _table_name, _df, _config):
        engine = create_engine(
            f"mysql+pymysql://{_config.USERNAME}:{_config.PASSWORD}@{_config.URL}:{_config.PORT}/{_database_name}")
        _df.to_sql(_table_name, con=engine, if_exists='append', index=False)


class JdbcConfigUtils:

    @classmethod
    def get_config_str(cls, _config: DbConfig, _table_name):
        # 连接数据库
        if isinstance(_config, MysqlConfig):
            return cls.get_mysql_config_str(_config, _table_name)
        else:
            raise Exception('不支持的数据库连接类型')

    @classmethod
    def get_mysql_config_str(cls, _config: MysqlConfig, _table_name):
        return f"""
                'connector' = '{_config.CONNECTOR}',
                'url' = 'jdbc:mysql://{_config.URL}:{_config.PORT}/{_config.DATABASE_NAME}',
                'table-name' = '{_table_name}',
                'username' = '{_config.USERNAME}',
                'password' = '{_config.PASSWORD}',
                'driver' = '{_config.DRIVER}'
        """

    @classmethod
    def get_connect(cls, _config: DbConfig):
        # 连接数据库
        if isinstance(_config, MysqlConfig):
            return cls.get_mysql_connect(_config)
        else:
            raise Exception('不支持的数据库连接类型')

    @classmethod
    def get_mysql_connect(cls, _config: MysqlConfig):
        # 连接数据库
        cnx = pymysql.connect(host=_config.URL,
                              port=_config.PORT,
                              user=_config.USERNAME,
                              password=_config.PASSWORD,
                              database=_config.DATABASE_NAME)
        return cnx

    @classmethod
    def execute_no_result_sql(cls, _config: DbConfig, _sql):
        with JdbcConfigUtils.get_connect(_config) as cnx:
            with cnx.cursor() as cursor:
                for _part_sql in _sql.split(';'):
                    _part_sql = _part_sql.strip()
                    if len(_part_sql) <= 0:
                        continue
                    cursor.execute(_part_sql)
                cnx.commit()

    @classmethod
    def get_table_data(cls, _config: DbConfig, _sql):
        with JdbcConfigUtils.get_connect(_config) as cnx:
            df = pd.read_sql(_sql, cnx)
            return df

    @classmethod
    def count_table(cls, _config: DbConfig, _database_name):
        _count_sql = f'SELECT COUNT(*) SIZE_NUM FROM {_database_name}'
        with JdbcConfigUtils.get_connect(_config) as cnx:
            df = pd.read_sql(_count_sql, cnx)
            return int(df.iloc[0]['SIZE_NUM'])

    @classmethod
    def get_table_data_with_chunk_size(cls, _config: DbConfig, _sql, _chunk_size):
        with JdbcConfigUtils.get_connect(_config) as cnx:
            df = pd.read_sql(_sql, cnx, chunksize=_chunk_size)
            return df

    @classmethod
    def insert_table_data(cls, _database_name, _table_name, _df, _config: MysqlConfig, _chunk_size=50000,
                          if_exists='append'):
        engine = create_engine(
            f"mysql+pymysql://{_config.USERNAME}:{_config.PASSWORD}@{_config.URL}:{_config.PORT}/{_database_name}")
        # 每次写入50000行数据
        _total_size = len(_df)
        if _total_size <= _chunk_size:
            _df.to_sql(_table_name, con=engine, if_exists='append', index=False)
        else:
            _start_index = 0
            _end_index = _chunk_size
            while _start_index < _total_size:
                _chunk = _df[_start_index:_end_index]
                _chunk.to_sql(_table_name, con=engine, if_exists=if_exists, index=False)
                _start_index += _chunk_size
                _end_index += _chunk_size

    @classmethod
    def insert_table_data_v2(cls, _database_name, _table_name, _df, _config: MysqlConfig,
                          _chunk_size=50000, if_exists='append'):
        """
        优化插入方法，使用LOAD DATA直接插入，数据库需要配配置才可以使用
        """
        temp_dir = tempfile.mkdtemp()
        conn = None
        try:
            # 创建数据库连接（关键：local_infile=True）
            conn = pymysql.connect(
                host=_config.URL,
                user=_config.USERNAME,
                password=_config.PASSWORD,
                port=_config.PORT,
                database=_database_name,
                local_infile=True  # 确保启用
            )

            _start_index = 0
            while _start_index < len(_df):
                _end_index = _start_index + _chunk_size
                _chunk = _df[_start_index:_end_index]

                # 生成临时CSV文件（Windows路径需转义）
                temp_csv = os.path.join(temp_dir, f"temp_{_start_index}.csv").replace('\\', '/')
                _chunk.to_csv(temp_csv, index=False, header=False, sep='\t')

                # 构建LOAD DATA命令
                load_sql = f"""
                    LOAD DATA LOCAL INFILE '{temp_csv}'
                    INTO TABLE {_table_name}
                    FIELDS TERMINATED BY '\\t'
                    LINES TERMINATED BY '\\n'
                """

                # 执行导入
                with conn.cursor() as cursor:
                    cursor.execute("SET autocommit = 0")
                    cursor.execute("SET unique_checks = 0")
                    cursor.execute("SET foreign_key_checks = 0")
                    cursor.execute(load_sql)
                    conn.commit()

                os.remove(temp_csv)
                _start_index = _end_index

        except Exception as e:
            if conn:
                conn.rollback()
            raise RuntimeError(f"数据插入失败: {str(e)}") from e
        finally:
            if conn:
                conn.close()
            if os.path.exists(temp_dir):
                shutil.rmtree(temp_dir)  # 使用shutil彻底删除目录

    @classmethod
    def check_database(cls, _database_name, _config: MysqlConfig):
        _sql = f"SHOW DATABASES LIKE '{_database_name}';"
        _df = cls.get_table_data(_config, _sql)
        return _df.shape[0] != 0

    @classmethod
    def drop_database(cls, _database_name, _config: MysqlConfig):
        _sql = f"DROP DATABASE {_database_name};"
        cls.execute_no_result_sql(_config, _sql)

    @classmethod
    def create_database(cls, _database_name, _config: MysqlConfig):
        _sql = f"CREATE DATABASE {_database_name} CHARACTER SET utf8mb4 COLLATE utf8mb4_general_ci;"
        cls.execute_no_result_sql(_config, _sql)

    @classmethod
    def delete_table_data(cls, _database_name, _table_name, _condition, _config):
        """
        通用删除方法，根据条件删除表中的数据
        :param _database_name: 数据库名称
        :param _table_name: 表名
        :param _condition: 删除条件（SQL WHERE 子句，例如 "city = '武汉市'"）
        :param _config: 数据库配置
        """
        # 构建 SQL 删除语句
        sql = f"DELETE FROM {_table_name} WHERE {_condition};"
        with cls.get_connect(_config) as cnx:
            with cnx.cursor() as cursor:
                cursor.execute(sql)
                cnx.commit()
                print(f"已删除表 {_table_name} 中满足条件的数据")

    @classmethod
    def delete_table_data_v2(cls, _database_name, _table_name, _condition, _config):
        """
        通用删除方法，根据条件分页删除表中的数据，每次删除10000条
        :param _database_name: 数据库名称
        :param _table_name: 表名
        :param _condition: 删除条件（SQL WHERE 子句，例如 "city = '武汉市'"）
        :param _config: 数据库配置
        """
        # 构建基础SQL语句（适用于MySQL，使用LIMIT分页删除）
        base_sql = f"DELETE FROM {_table_name} WHERE {_condition} LIMIT 50000;"
        total_deleted = 0
        print(base_sql)
        with cls.get_connect(_config) as cnx:
            with cnx.cursor() as cursor:
                while True:
                    cursor.execute(base_sql)
                    deleted_rows = cursor.rowcount
                    cnx.commit()  # 提交当前事务
                    total_deleted += deleted_rows
                    print(f"本次删除 {deleted_rows} 条数据，累计删除 {total_deleted} 条")

                    # 如果删除行数小于10000，表示没有更多数据
                    if deleted_rows < 1000:
                        break

        print(f"已删除表 {_table_name} 中满足条件的数据，共删除 {total_deleted} 条")

    @classmethod
    def truncate_table_data(cls, _database_name, _table_name, _config):
        """
        通用删除方法，根据条件删除表中的数据
        :param _database_name: 数据库名称
        :param _table_name: 表名
        :param _config: 数据库配置
        """
        # 构建 SQL 删除语句
        sql = f"TRUNCATE TABLE `{_table_name}`;"
        with cls.get_connect(_config) as cnx:
            with cnx.cursor() as cursor:
                cursor.execute(sql)
                cnx.commit()

    @staticmethod
    def generate_uuid():
        return str(uuid.uuid4()).replace('-', '')
